# -*- coding: utf-8 -*-

import random
import os
import numpy as np
import cv2
import scipy.io
import h5py
import tensorflow as tf
from mcnn import cnn_column_1
from read_data import get_images_path, process_image


# def read_image_and_ground_truth():
#     folder = "/home/lijun/data/dataset/DataSet/ShanghaiTech/part_A_final/train_data/images"
#     images_name = [name for name in os.listdir(folder) if "jpg" in name]
#     # name = random.choice(images_name)
#     name = images_name[0]
#     image_path = os.path.join(folder, name)
#     mat_path = image_path.replace("images/", "ground_truth/GT_").replace("jpg", "mat")
#     mat = scipy.io.loadmat(mat_path)
#     locations = mat["image_info"][0][0][0][0][0]
#     print(locations.shape)
#     image = cv2.imread(image_path)
#     cv2.imshow("show", image)
#     image = np.expand_dims(image, axis=0)
#     return image


def read_image_and_ground_truth():
    folder = "/home/lijun/Dataset/crowd_counting/part_b/train_multi_1000/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    name = random.choice(images_name)
    image_path = os.path.join(folder, name)
    image = cv2.imread(image_path)
    cv2.imshow("image", image)
    image = image.astype(dtype=np.float32, copy=False)
    image = np.divide(image, 255)

    mat_path = image_path.replace("images", "ground_truth").replace("jpg", "h5")
    ground_truth_file = h5py.File(mat_path, "r")
    ground_truth = np.array(ground_truth_file["data"])
    ground_truth_file.close()

    image = np.expand_dims(image, axis=0)
    return image, ground_truth


def show_ground_truth(ground_truth, name="color_map"):
    min_value, max_value = np.min(ground_truth), np.max(ground_truth)
    img = ((ground_truth - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.resize(img, (img.shape[1]*4, img.shape[0]*4))
    color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    cv2.imshow(name, color_map)


def main():
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32, shape=[1, 192, 256, 3])
        inference_op = cnn_column_1(in_put=input_x)
        saver = tf.train.Saver()
        with tf.Session() as session:
            saver.restore(session, "./model/cnn_column_1-40000")
            image, ground_truth = read_image_and_ground_truth()
            inference = session.run(inference_op, feed_dict={input_x: image})
            inference = inference[0][:, :, 0]
            print(np.sum(inference)/1000, np.sum(ground_truth)/1000)
            show_ground_truth(inference, "inference")
            show_ground_truth(ground_truth, "ground_truth")
            cv2.waitKey()
            cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
